#pragma once

#include "limbo/details/checksum.h"
#include "limbo/errc.h"
#include "limbo/ip/state.hpp"
#include "limbo/result.h"
#include "opts.h"
#include "udp_context.hpp"
#include <cstdint>

namespace limbo::udp {

template <typename State> struct UdpPacket {
  using LowerState = typename State::LowerState;
  using Context = UdpContext<State>;
  using Container = typename LowerState::Packet;
  static constexpr uint8_t SZ = 8;

  struct PseudoHeader {
    uint32_t src;
    uint32_t dst;
    uint8_t zeroes;
    uint8_t proto;
    uint16_t length;
    uint16_t checksum;
  };

  uint16_t source_port;
  uint16_t dest_port;
  uint16_t length;
  uint16_t checksum;
  Chunk payload;
  Container *container;

  inline static uint32_t prefix(const Context &) { return SZ; }

  static Result send(State &state, const Context &ctx, Chunk buff,
                     const Chunk payload) {
    if (buff.size() < payload.size() + SZ) {
      auto code = (uint32_t)Errc::no_enough_space;
      return Result::make_error(code, &state);
    }

    alignas(4) char data[SZ];
    auto ptr = (unsigned char *)data;

    *((uint16_t *)ptr) = order::native_to_big(ctx.source_port);
    ptr += sizeof(uint16_t);

    *((uint16_t *)ptr) = order::native_to_big(ctx.dest_port);
    ptr += sizeof(uint16_t);

    auto sz = (uint16_t)(payload.size() + SZ);
    auto sz_be = order::native_to_big(sz);
    *((uint16_t *)ptr) = sz_be;
    ptr += sizeof(uint16_t);

    /* checksum */
    auto checksum = (uint16_t *)ptr;
    *checksum = 0;
    ptr += sizeof(uint16_t);

    memcpy(buff.data(), data, SZ);

    auto consumed = Chunk(buff.data(), SZ);
    consumed << payload;
    if constexpr (State::opts & Opts::calculate_checksum) {
      alignas(4) PseudoHeader h;
      h.src = ctx.ip_context->source.as_net();
      h.dst = ctx.ip_context->destination.as_net();
      h.zeroes = 0;
      h.proto = (uint8_t)ip::Proto::udp;
      h.length = sz_be;
      h.checksum = 0;
      auto p_cs = details::checksum(&h, sizeof(h));
      auto h_cs = details::checksum(consumed.data(), consumed.size());
      auto cs = details::checksum_add(p_cs, h_cs);
      auto cs_ptr = (uint16_t *)(buff.data() + 6);
      *cs_ptr = order::native_to_big((uint16_t)~cs);
    }

    return Result::make_consumed(&state, consumed);
  }

  Result recv(Chunk buff, State &state, LowerState *ip_state) {
    auto sz = buff.size();
    auto ptr = (unsigned char *)buff.data();
    if (sz < SZ) {
      return Result::make_demand(&state, SZ - sz);
    }
    auto start = ptr;

    uint16_t source_port;
    memcpy(&source_port, ptr, sizeof(source_port));
    ptr += sizeof(source_port);
    this->source_port = order::big_to_native(source_port);

    uint16_t dest_port;
    memcpy(&dest_port, ptr, sizeof(dest_port));
    ptr += sizeof(dest_port);
    this->dest_port = order::big_to_native(dest_port);

    uint16_t length;
    memcpy(&length, ptr, sizeof(length));
    ptr += sizeof(length);
    this->length = order::big_to_native(length);
    if (sz < this->length) {
      return Result::make_demand(&state, this->length - sz);
    }

    uint16_t checksum;
    memcpy(&checksum, ptr, sizeof(checksum));
    ptr += sizeof(checksum);
    this->checksum = order::big_to_native(checksum);

    payload = Chunk(ptr, this->length - SZ);
    auto consumed = Chunk(start, this->length);
    container = &ip_state->get_parsed();

    if constexpr (State::opts & Opts::validate_checksum) {
      alignas(4) PseudoHeader h;
      h.src = container->source.as_net();
      h.dst = container->destination.as_net();
      h.zeroes = 0;
      h.proto = (uint8_t)ip::Proto::udp;
      h.length = length;
      h.checksum = 0;
      auto p_cs = details::checksum(&h, sizeof(h));
      auto h_cs = details::checksum(start, this->length);
      auto cs_1 = details::checksum_add(p_cs, h_cs);
      auto cs_2 = details::checksum_sub(cs_1, this->checksum);
      auto cs = (uint16_t)~cs_2;
      if (cs != this->checksum) {
        auto code = (uint32_t)Errc::udp_checksum_mismatch;
        return Result::make_error(code, &state, consumed);
      }
    }

    /* success */
    return Result::make_consumed(&state, consumed);
  }
};

} // namespace limbo::udp
